# -*- coding: utf-8 -*-
#!/usr/bin/env python
import rospy
import numpy as np
from sensor_msgs.msg import PointCloud2
from nav_msgs.msg import Odometry
import sensor_msgs.point_cloud2 as pc2
from tf.transformations import quaternion_matrix

uav_id = $id

class PointCloudFilter:
    def __init__(self):
        rospy.init_node('pointcloud_filter',anonymous=True)

        # 订阅原始点云和里程计
        self.odom_sub = rospy.Subscriber(
            '/iris_{}/transfer/Odometry'.format(uav_id), 
            Odometry, 
            self.odom_callback
        )
        self.pc_sub = rospy.Subscriber(
            '/iris_{}/transfer/cloud_registered'.format(uav_id),
            PointCloud2,
            self.pc_callback
        )

        # 发布筛选后的点云
        self.filtered_pub = rospy.Publisher(
            '/iris_{}/fov/cloud'.format(uav_id),
            PointCloud2,
            queue_size=10
        )

        # 更新里程计时间戳
        self.odom_pub = rospy.Publisher(
            '/iris_{}/fov/Odometry'.format(uav_id),
            Odometry,
            queue_size=10
        )

        self.current_pose = None

    def odom_callback(self, msg):
        """更新当前无人机位姿"""
        self.current_pose = msg.pose.pose
        msg.header.stamp = rospy.Time.now()
        self.odom_pub.publish(msg)

    def pc_callback(self, msg):
        """处理点云数据"""
        if self.current_pose is None:
            rospy.logwarn_once("等待初始化：未收到里程计数据")
            return

        # 获取位姿参数
        pos = self.current_pose.position
        ori = self.current_pose.orientation
        
        # 关键修正：调整旋转矩阵应用方式
        q = [ori.x, ori.y, ori.z, ori.w]
        R = quaternion_matrix(q)[:3, :3]  # 世界坐标系到机体坐标系的旋转矩阵
        
        # 读取点云数据
        points = np.array(list(pc2.read_points(
            msg, 
            field_names=("x", "y", "z"),
            skip_nans=True
        )))
        
        if points.size == 0:
            return

        # 坐标转换流程 (修正方向)
        translated = points - [pos.x, pos.y, pos.z]
        body_coords = np.dot(translated, R)  # 关键修改：使用R而非R.T
        
        # 计算方位角 (修正参数顺序)
        angles_deg = np.degrees(np.arctan2(body_coords[:,1], body_coords[:,0]))
        
        # 创建筛选掩码
        angle_mask = np.abs(angles_deg) <= 43
        filtered_points = points[angle_mask]

        # 构造输出点云
        header = msg.header
        header.stamp = rospy.Time.now()
        filtered_cloud = pc2.create_cloud_xyz32(header, filtered_points)
        self.filtered_pub.publish(filtered_cloud)

if __name__ == '__main__':
    PointCloudFilter()
    rospy.spin()